# -*- coding: utf-8 -*-
import cv2
import os
from tqdm import tqdm
import random
import shutil
import numpy as np
import matplotlib.pyplot as plt
import skimage

def get_label(tmp_subfolder):
        label_dic = {'HCC': 3, 'CYST': 0, 'FNH': 1, 'HA': 2, 'ICC': 4, 'META': 5, 'Hemangioma': 2, 'nodule': 0}
        if 'HCC' in tmp_subfolder:
            return label_dic['HCC']
        elif 'CYST' in tmp_subfolder:
            return label_dic['CYST']
        elif 'FNH' in tmp_subfolder:
            return label_dic['FNH']
        elif 'HA' in tmp_subfolder:
            return label_dic['HA']
        elif 'ICC' in tmp_subfolder:
            return label_dic['ICC']
        elif 'META' in tmp_subfolder:
            return label_dic['META']
        elif 'Hemangioma' in tmp_subfolder:
            return label_dic['Hemangioma']
        elif 'nodule' in tmp_subfolder:
            return label_dic['nodule']


def process_images(input_folder, output_root_folder, train_ratio):
    # 确保保存分割图像的文件夹存在
    if not os.path.exists(output_root_folder):
        os.makedirs(output_root_folder)
    # 遍历文件夹，一个文件夹对应一个类别
    classes = [cla for cla in os.listdir(input_folder) if os.path.isdir(os.path.join(input_folder, cla))]
    # 排序，保证顺序一致
    classes.sort()
    # 遍历每个文件夹下的文件
    images_path = []
    label = []
    for cla in tqdm(classes):
        cla_path = os.path.join(input_folder, cla)
        images = os.listdir(cla_path)  # 获取该类别下的所有文件名
        if '.DS_Store' in images:
            images.remove('.DS_Store')
        for patient in images:
            patient_path = os.path.join(cla_path,patient)
            images_path.append(patient_path)
            label.append(get_label(patient_path))

    for i, img in enumerate(images_path):
        img_path2 = os.path.join(img, 'a.npz')
        img_path3 = os.path.join(img, 'p.npz')
        img_path4 = os.path.join(img, 'v.npz')
        data2 = np.load(img_path2)
        data3 = np.load(img_path3)
        data4 = np.load(img_path4)
        image_a = data2['image']
        mask_a = data2['mask']
        image_p = data3['image']
        mask_p = data3['mask']
        image_v = data4['image']
        mask_v = data4['mask']
        z_a, x_a, y_a = np.where(mask_a > 0)
        z_a_start = np.min(z_a)
        z_a_end = np.max(z_a)
        x_a_start = np.min(x_a)
        x_a_end = np.max(x_a)
        y_a_start = np.min(y_a)
        y_a_end = np.max(y_a)

        z_p, x_p, y_p = np.where(mask_p > 0)
        z_p_start = np.min(z_p)
        z_p_end = np.max(z_p)
        x_p_start = np.min(x_p)
        x_p_end = np.max(x_p)
        y_p_start = np.min(y_p)
        y_p_end = np.max(y_p)

        z_v, x_v, y_v = np.where(mask_v > 0)
        z_v_start = np.min(z_v)
        z_v_end = np.max(z_v)
        x_v_start = np.min(x_v)
        x_v_end = np.max(x_v)
        y_v_start = np.min(y_v)
        y_v_end = np.max(y_v)

        tumor_a = image_a[z_a_start:z_a_end+1, x_a_start:x_a_end+1, y_a_start:y_a_end+1]
        tumor_p = image_p[z_p_start:z_p_end+1, x_p_start:x_p_end+1, y_p_start:y_p_end+1]
        tumor_v = image_v[z_v_start:z_v_end+1, x_v_start:x_v_end+1, y_v_start:y_v_end+1]

        transformed_tumor_a = skimage.transform.resize(tumor_a, output_shape=(64, 64, 64), anti_aliasing=None, order=1, preserve_range=True).astype(np.uint8)
        transformed_tumor_p = skimage.transform.resize(tumor_p, output_shape=(64, 64, 64), anti_aliasing=None, order=1, preserve_range=True).astype(np.uint8)
        transformed_tumor_v = skimage.transform.resize(tumor_v, output_shape=(64, 64, 64), anti_aliasing=None, order=1, preserve_range=True).astype(np.uint8)
        if label[i]==0:
            name_label = "CYST"
        elif label[i]==1:
            name_label = "FNH"
        elif label[i]==2:
            name_label = "HA"
        elif label[i]==3:
            name_label = "HCC"
        elif label[i]==4:
            name_label = "ICC"
        elif label[i]==5:
            name_label = "META"

        t_path = os.path.join(output_root_folder, name_label)
        output_patient_folder = os.path.join(t_path,f'patient_{i}')
        if not os.path.exists(output_patient_folder):
            os.makedirs(output_patient_folder)

        np.savez(os.path.join(output_patient_folder, 'a.npz'), image=transformed_tumor_a)
        np.savez(os.path.join(output_patient_folder, 'p.npz'), image=transformed_tumor_p)
        np.savez(os.path.join(output_patient_folder, 'v.npz'), image=transformed_tumor_v)
    # # 遍历原始图像文件夹中的所有文件夹，假设每个子文件夹包含一个类别的图像
    # for class_folder_name in os.listdir(input_folder):
    #     class_folder_path = os.path.join(input_folder, class_folder_name)
    #     if not os.path.isdir(class_folder_path):
    #         continue
    #
    #     # 获取类别名称
    #     class_name = class_folder_name
    #
    #     # 创建保存分割图像的文件夹路径
    #     output_class_folder = os.path.join(output_root_folder, class_name)
    #     if not os.path.exists(output_class_folder):
    #         os.makedirs(output_class_folder)
    #
    #     # 记录是否已经处理了带有DMK后缀的图片
    #     dmk_processed = False
    #
    #     # 遍历当前类别文件夹中的所有图像文件
    #     for filename in os.listdir(class_folder_path):
    #
    #         if filename.endswith(".jpg") or filename.endswith(".png"):  # 仅处理 JPG 或 PNG 格式的图像
    #             if "DMK" in filename:
    #                 # 如果有DMK后缀的图片存在，并且还未处理过，则使用这张图片作为原图
    #                 if not dmk_processed:
    #                     original_image_path = os.path.join(class_folder_path, filename)
    #                     dmk_processed = True
    #                 else:
    #                     continue  # 如果已经处理了带有DMK后缀的图片，则跳过当前图片
    #             else:
    #                 # 如果没有后缀的图片，或者有DMK后缀的图片不存在，则使用当前图片作为原图
    #                 original_image_path = os.path.join(class_folder_path, filename)
    #
    #             # 加载原始图像
    #             original_image = cv2.imread(original_image_path)
    #             print("oriinal_image_path:",original_image_path)
    #             if original_image is None:
    #                 print(f"Error: Unable to load image '{original_image_path}'")
    #                 continue
    #
    #             # 遍历可能存在的ROI文件
    #
    #             for roi_suffix in ["_ROI1"]:
    #             #for roi_suffix in ["_ROI1", "_ROI2", "_ROI3"]:
    #                 n=0
    #                 roi_image_filename = filename.replace(".png", f"{roi_suffix}.png")
    #                 roi_image_path = os.path.join(class_folder_path, roi_image_filename)
    #                 print("roi_image_path",roi_image_path)
    #                 roi_image = cv2.imread(roi_image_path, cv2.IMREAD_GRAYSCALE) if os.path.exists(roi_image_path) else None
    #
    #                 # 如果ROI图像存在，则进行分割处理
    #                 if roi_image is not None:
    #                     #print(roi_image.shape)
    #                     #plt.figure
    #                     #plt.imshow(original_image, cmap='gray')
    #                     #plt.contour(roi_image[:,:,0])
    #                     #plt.show()
    #
    #                     x_, y_ = np.where(roi_image[:,:]>0)
    #
    #                     x_start = np.min(x_)
    #                     x_end = np.max(x_)
    #
    #                     y_start = np.min(y_)
    #                     y_end = np.max(y_)
    #
    #                     original_image = original_image[x_start: x_end+1, y_start:y_end+1,:]
    #                     roi_image = roi_image[x_start: x_end+1, y_start:y_end+1]
    #
    #                     #plt.figure
    #                     #plt.imshow(original_image, cmap='gray')
    #                     #plt.contour(roi_image[:,:,0])
    #                     #plt.show()
    #
    #                     output_image_path = os.path.join(output_class_folder, f"segmented_{roi_image_filename}")
    #                     cv2.imwrite(output_image_path,original_image)
    #                     n+=1
    #                     print(f"Segmented image saved: '{output_image_path}'")
    #
    #                     '''segmented_image = cv2.bitwise_and(original_image, original_image, mask=roi_image)
    #                     output_image_path = os.path.join(output_class_folder, f"segmented_{roi_image_filename}")
    #                     cv2.imwrite(output_image_path, segmented_image)
    #                     print(n,f"Segmented image saved: '{output_image_path}'")'''
    #                 # 如果ROI图像不存在，并且是第一个ROI，则输出找不到ROI的信息
    #                 elif roi_suffix == "_ROI1":
    #                     print(f"No ROI found for '{filename}' in '{class_folder_path}'")
    #
    # print(n,"Processing completed.")

    #     # 分割训练集和测试集
    # for class_name in os.listdir(output_root_folder):
    #     class_folder = os.path.join(output_root_folder, class_name)
    #     if not os.path.isdir(class_folder):
    #         continue

    #     train_folder = os.path.join(output_root_folder, "train", class_name)
    #     test_folder = os.path.join(output_root_folder, "test", class_name)

    #     if not os.path.exists(train_folder):
    #         os.makedirs(train_folder)
    #     if not os.path.exists(test_folder):
    #         os.makedirs(test_folder)

    #     # 遍历当前类别文件夹中的所有图像文件
    #     images = os.listdir(class_folder)
    #     random.shuffle(images)

    #     train_count = int(len(images) * train_ratio)
    #     train_images = images[:train_count]
    #     test_images = images[train_count:]

    #     for image in train_images:
    #         shutil.move(os.path.join(class_folder, image), os.path.join(train_folder, image))

    #     for image in test_images:
    #         shutil.move(os.path.join(class_folder, image), os.path.join(test_folder, image))

    # print("Splitting into train and test completed.")

if __name__ == '__main__':
    # 原始图像文件夹路径和保存分割图像的文件夹路径
    input_folder = "/home/hzt/whh/pycharm_project_179/data"
    output_root_folder = "/home/hzt/whh/pycharm_project_179/output"
    # input_folder = "C:/Users/ChinHaiMei/OneDrive/Desktop/test/input"
    # output_root_folder = "C:/Users/ChinHaiMei/OneDrive/Desktop/test/output"
    train_ratio = 0.8  # 训练集所占比例
    process_images(input_folder, output_root_folder, train_ratio)
